#include <assert.h>
#include <stdio.h>
#include <QFile>

#include "pdbfileplugin.h"
#include "../../data/proteindata.h"

using namespace PBVP;

PDBFilePlugin::PDBFilePlugin()
{
	curChain = '\0';
	curResId = 0;
}

PDBFilePlugin::~PDBFilePlugin()
{
}

void PDBFilePlugin::initialize(PBVP::MainInterface *pMainInterface)
{
	// nothing here
}

DataInterface* PDBFilePlugin::loadFile(const QString &fname)
{
	QFile file(fname);
	if (!file.open(QIODevice::ReadOnly))
	{
		return NULL;
	}
	ProteinData *pPDBData = new ProteinData();
	PDBLineType recordType;
	do 
	{
		QByteArray line;
		recordType = readPDBLine(file, line);
		// printf("%s\n", line.data());
		switch(recordType) {
		case PDB_HEADER:
			parseHeader(pPDBData, line);
			break;
		case PDB_COMPND:
			parseCompound(pPDBData, line);
			break;
		case PDB_SEQRES:
			parseResidues(pPDBData, line);
			break;
		case PDB_ATOM:
			parseAtoms(pPDBData, line);
			break;
		case PDB_HETATM:		// non-polymer atoms
			parseHetAtoms(pPDBData, line);
			break;
		case PDB_HELIX:
			parseHelix(pPDBData, line);
			break;
		case PDB_SHEET:
			parseSheet(pPDBData, line);
			break;
		default:
			break;
		}
	} while(recordType != PDB_END && recordType != PDB_EOF);
	file.close();
#ifdef _DEBUG
	pPDBData->print();
#endif
	return pPDBData;
}

bool PDBFilePlugin::saveFile(DataInterface *pData, const QString &fname)
{
	// Don't support saving as PDB file
	return false;
}

//////////////////////////////////////////////////////////////////////////
// Private functions
//////////////////////////////////////////////////////////////////////////
PDBFilePlugin::PDBLineType PDBFilePlugin::readPDBLine(QFile& file, QByteArray& line)
{
	line = file.readLine();
	if (line.length() == 0)
	{ 
		return PDB_EOF;
	}
	line = line.trimmed();
	if (line.startsWith("END")) return PDB_END;
	if (line.startsWith("ATOM")) return PDB_ATOM;
	if (line.startsWith("HETATM")) return PDB_HETATM;
	if (line.startsWith("HEADER")) return PDB_HEADER;
	if (line.startsWith("REMARK")) return PDB_REMARK;
	if (line.startsWith("CONECT")) return PDB_CONECT;
	if (line.startsWith("CRYST1")) return PDB_CRYST1;
	if (line.startsWith("SHEET")) return PDB_SHEET;
	if (line.startsWith("TURN")) return	PDB_TURN;
	if (line.startsWith("SEQRES")) return PDB_SEQRES;
	if (line.startsWith("HET   ")) return PDB_HET;
	if (line.startsWith("HELIX")) return PDB_HELIX;
	if (line.startsWith("COMPND")) return PDB_COMPND;

	// Treat the rest as unknown
	return PDB_UNKNOWN;
}

void PDBFilePlugin::parseHeader(ProteinData *pProtein, const QByteArray &line)
{
	ProteinMetaData meta;
	meta.classification = QString(subArray(line, 10, 40));
	meta.date = QString(subArray(line, 50, 9));
	meta.ID = QString(subArray(line, 62, 4));

	pProtein->setMetaData(meta);
}

void PDBFilePlugin::parseCompound(ProteinData *pProtein, const QByteArray &line)
{
	// Do not do anything here yet
}

void PDBFilePlugin::parseResidues(ProteinData *pProtein, QByteArray &line)
{
	char chain = line[11];
	// int nRes = subArray(line, 13, 4).toInt();
	if(chain != curChain) { // start of a new chain
		curChain = chain;
		pProtein->addChain(chain);
	} 
}

void PDBFilePlugin::parseAtoms(PBVP::ProteinData *pProtein, QByteArray &line)
{
	int atomId = subArray(line, 6, 5).toInt();
	char chain = subArray(line, 21, 1)[0];
	
	QString atomName(subArray(line, 12, 4));
	// the index of the residue in the chain
	int resId = subArray(line, 22, 4).toInt();
	QString resName(subArray(line, 17, 3));

	if(atomId == 1 || chain != curChain ) { // the first atom in the chain
		curChain = chain;
		// add a dummy atom for the end of last chain
		// artifact of PDB format
		/*if(atomId > 1) {
			pChain->addAtom(Atom());
		}*/
		pChain = pProtein->getChain(chain);
		// chain residue usually numbered starting at 1, but there are exceptions
		pChain->setChainOffset(resId);
		pChain->addResidue(resName);
		curResId = resId;
		
	} else if(resId != curResId) {	// not new chain but new residue
		curResId = resId;
		pChain->addResidue(resName);
	}

	float x = subArray(line, 30, 8).toFloat();
	float y = subArray(line, 38, 8).toFloat();
	float z = subArray(line, 46, 8).toFloat();
	float occupancy = subArray(line, 54, 6).toFloat();
	float tempFactor = subArray(line, 60, 6).toFloat();

	Atom atom(atomName, resName);
	atom.position = Vector3(x, y, z);
	atom.chainID = chain;
	atom.residueIndex = resId;
	pChain->addAtom(atom);
}

void PDBFilePlugin::parseHetAtoms(ProteinData *pProtein, QByteArray &line)
{
	QString ligName(subArray(line, 17, 3));
	char chainName = line[21];
	int resId = subArray(line, 22, 4).toInt();

	if(ligName == "HOH") return;	// ignore H2O atoms
	if(ligName != curLigand || chainName != curChain || resId != curResId) { // first atom in the ligand
		curLigand = ligName;
		curChain = chainName;
		curResId = resId;
		pLigand = pProtein->addLigand(ligName);
		if (chainName != ' ')
		{
			pChain = pProtein->getChain(chainName);
		}
	}
	// pLigand->incrementAtomCount();
	QString atomName(subArray(line, 76, 2));
	float x = subArray(line, 30, 8).toFloat();
	float y = subArray(line, 38, 8).toFloat();
	float z = subArray(line, 46, 8).toFloat();
	float occupancy = subArray(line, 54, 6).toFloat();
	float tempFactor = subArray(line, 60, 6).toFloat();

	Atom atom(atomName);
	atom.position = Vector3(x, y, z);
	pLigand->addAtom(atom);
}

void PDBFilePlugin::parseHelix(ProteinData *pProtein, QByteArray &line)
{
	// TODO
}

void PDBFilePlugin::parseSheet(ProteinData *pProtein, QByteArray &line)
{
	// TODO
}

QByteArray PDBFilePlugin::subArray(const QByteArray& line, int start, int len)
{
	assert(start+len <= line.length());
	const char* ptr = line.data();
	QByteArray sub(ptr+start, len);
	return sub;
}

Q_EXPORT_PLUGIN2(pdbfile, PDBFilePlugin)